package datastream.api.operator.group;

import datastream.api.operator.group.pojo.ReduceModel;
import datastream.api.operator.group.source.ReduceSource;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStreamSink;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

public class Reduce2 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);

        DataStreamSource<ReduceModel> source = env.addSource(new ReduceSource());

        SingleOutputStreamOperator<ReduceModel> map = source.map(new MapFunction<ReduceModel, ReduceModel>() {
            @Override
            public ReduceModel map(ReduceModel value) throws Exception {
                value.setCount(1);
                value.setAvgIncome(value.getIncome());
                return value;
            }
        });

        KeyedStream<ReduceModel, Integer> keyBy = map.keyBy(new KeySelector<ReduceModel, Integer>() {
            @Override
            public Integer getKey(ReduceModel value) throws Exception {
                return value.getProduceId();
            }
        });

        SingleOutputStreamOperator<ReduceModel> reduce = keyBy.reduce(new ReduceFunction<ReduceModel>() {
            @Override
            public ReduceModel reduce(ReduceModel value1, ReduceModel value2) throws Exception {

                int count = value1.getCount() + value2.getCount();
                double totalIncome = value1.getIncome() + value2.getIncome();
                value1.setIncome(totalIncome);
                value1.setCount(count);
                value1.setAvgIncome(totalIncome / count);

                // 与 Reduce1 的结果一致
                return ReduceModel.builder()
                        .produceId(value1.getProduceId())
                        .income(totalIncome)
                        .count(count)
                        .avgIncome(totalIncome / count)
                        .build();
            }
        });

        DataStreamSink<ReduceModel> sink = reduce.print();

        env.execute();
    }
}
